from py2neo import Graph
import logging
import time
import sys

# Neo4j 连接配置
CONFIG = {
    "uri": "bolt://localhost:7687",  # Neo4j 服务器地址
    "auth": ("neo4j", "12345678")  # 认证信息 (用户名, 密码)
}


class Neo4jDeepCleaner:
    def __init__(self, uri: str, username: str, password: str):
        """初始化数据库清理器
        Args:
            uri: Neo4j 连接地址
            username: 数据库用户名
            password: 数据库密码
        """
        self.graph = Graph(uri, auth=(username, password))
        self._setup_logging()
        self.logger = logging.getLogger(__name__)
        self.logger.info("Neo4j深度清理器初始化完成")

    def _setup_logging(self):
        """配置日志系统"""
        logging.basicConfig(
            level=logging.INFO,  # 日志级别
            format='%(asctime)s [%(levelname)s] %(message)s',  # 日志格式
            handlers=[logging.StreamHandler()]  # 输出到控制台
        )

    def execute_deep_clean(self):
        """执行深度清理流程"""
        start_time = time.time()
        try:
            self.logger.info("开始深度清理数据库...")

            # 第一阶段：删除所有节点和关系
            self._delete_all_data()

            # 第二阶段：清理约束和索引
            self._clean_constraints_and_indexes()

            # 第三阶段：重置属性键
            self._reset_property_keys()

            duration = time.time() - start_time
            self.logger.info(f"✅ 深度清理完成 | 总耗时: {duration:.2f}秒")

        except Exception as e:
            self.logger.error(f"❌ 清理失败: {str(e)}")
            raise

    def _delete_all_data(self):
        """删除所有节点和关系"""
        self.logger.info("正在删除所有数据...")
        start_time = time.time()

        # 使用分批事务处理避免内存溢出
        self.graph.run("""
        MATCH (n)
        CALL { WITH n DETACH DELETE n }
        IN TRANSACTIONS OF 10000 ROWS
        """)

        duration = time.time() - start_time
        self.logger.info(f"数据删除完成 | 耗时: {duration:.2f}秒")

    def _clean_constraints_and_indexes(self):
        """清理所有约束和索引"""
        self.logger.info("正在清理约束和索引...")
        start_time = time.time()

        # 删除约束
        constraint_count = 0
        for constraint in self.graph.run("SHOW CONSTRAINTS").data():
            self.graph.run(f"DROP CONSTRAINT {constraint['name']}")
            constraint_count += 1
            self.logger.debug(f"已删除约束: {constraint['name']}")

        # 删除索引（排除系统索引）
        index_count = 0
        for index in self.graph.run("SHOW INDEXES WHERE type <> 'LOOKUP'").data():
            self.graph.run(f"DROP INDEX {index['name']}")
            index_count += 1
            self.logger.debug(f"已删除索引: {index['name']}")

        duration = time.time() - start_time
        self.logger.info(f"已删除 {constraint_count} 个约束和 {index_count} 个索引 | 耗时: {duration:.2f}秒")

    def _reset_property_keys(self):
        """重置属性键（需要APOC插件支持）"""
        self.logger.info("正在重置属性键...")
        start_time = time.time()

        try:
            # 获取当前属性键数量
            original_count = self.graph.run("CALL db.propertyKeys() YIELD propertyKey RETURN count(*)").data()[0][
                'count(*)']

            # 使用APOC清理
            self.graph.run("CALL apoc.schema.assert({}, {})")

            # 验证清理结果
            new_count = self.graph.run("CALL db.propertyKeys() YIELD propertyKey RETURN count(*)").data()[0]['count(*)']

            duration = time.time() - start_time
            if new_count < original_count:
                self.logger.info(f"属性键从 {original_count} 个减少到 {new_count} 个 | 耗时: {duration:.2f}秒")
            else:
                self.logger.warning("属性键数量未变化（可能需要重启数据库）")

        except Exception as e:
            self.logger.warning(f"APOC清理失败（社区版限制）: {str(e)}")
            self.logger.info("提示: 重启Neo4j服务可完全重置属性键")


if __name__ == "__main__":
    try:
        cleaner = Neo4jDeepCleaner(CONFIG["uri"], CONFIG["auth"][0], CONFIG["auth"][1])
        cleaner.execute_deep_clean()

    except Exception as e:
        logging.error(f"❌ 致命错误: {str(e)}", exc_info=True)
        sys.exit(1)